"""
股票基本信息数据抓取模块
"""
import akshare as ak
import tushare as ts
import pandas as pd
import numpy as np
from loguru import logger
from sqlalchemy.orm import Session

from config.settings import TUSHARE_TOKEN
from models.stock_models import StockBasicInfo
from database.db_engine import SessionLocal

# 设置tushare接口token
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()

def fetch_stock_list():
    """
    获取A股上证所有股票列表
    
    Returns:
        pandas.DataFrame: 股票列表数据
    """
    try:
        logger.info("开始获取A股上证股票列表...")
        
        # 通过akshare获取股票列表
        stock_info_sh_df = ak.stock_info_sh_name_code()
        logger.info(f"成功获取上证股票列表，共 {len(stock_info_sh_df)} 条记录")
        
        # 转换列名和格式
        stock_info_sh_df = stock_info_sh_df.rename(
            columns={
                'code': 'stock_code',
                'name': 'name'
            }
        )
        
        # 添加交易所标识
        stock_info_sh_df['ts_code'] = stock_info_sh_df['stock_code'] + '.SH'
        
        return stock_info_sh_df
    except Exception as e:
        logger.error(f"获取股票列表失败: {e}")
        return pd.DataFrame()

def fetch_stock_basic_info(stock_code_list=None):
    """
    获取股票的基本信息（市值、市盈率等）
    
    Args:
        stock_code_list (list, optional): 股票代码列表. 默认为None，表示获取所有股票信息
    
    Returns:
        pandas.DataFrame: 股票基本信息数据
    """
    try:
        logger.info("开始获取股票基本信息...")
        
        # 通过tushare获取股票基本信息
        basic_info = pro.stock_basic(
            exchange='',
            list_status='L',  # 上市状态: L上市 D退市 P暂停上市
            fields='ts_code,symbol,name,area,industry,market,list_date,total_share,circulating_share,list_status'
        )
        
        # 如果指定了股票代码列表，则只获取指定股票的信息
        if stock_code_list:
            # 将股票代码转换为ts_code格式（添加.SH或.SZ后缀）
            ts_code_list = []
            for code in stock_code_list:
                if code.startswith('6'):
                    ts_code_list.append(f"{code}.SH")
                elif code.startswith(('0', '3')):
                    ts_code_list.append(f"{code}.SZ")
            
            basic_info = basic_info[basic_info['ts_code'].isin(ts_code_list)]
            logger.info(f"已筛选指定的 {len(stock_code_list)} 只股票信息")
        
        # 获取当前最新行情数据以计算市值、PE、PB等
        quotes = pro.daily_basic(
            trade_date=pd.Timestamp.now().strftime('%Y%m%d'),
            fields='ts_code,total_mv,circ_mv,pe,pe_ttm,pb'
        )
        
        # 合并数据
        result = pd.merge(basic_info, quotes, on='ts_code', how='left')
        
        # 重命名列
        result = result.rename(
            columns={
                'symbol': 'stock_code',
                'total_mv': 'total_capital',
                'circ_mv': 'circulating_capital',
                'list_status': 'status',
                'pe_ttm': 'pe'  # 使用滚动市盈率（TTM）作为PE值
            }
        )
        
        # 将nan值转换为None
        result = result.replace({np.nan: None})
        
        logger.info(f"成功获取股票基本信息，共 {len(result)} 条记录")
        return result
    
    except Exception as e:
        logger.error(f"获取股票基本信息失败: {e}")
        return pd.DataFrame()

def update_stock_basic_info(stock_code_list=None):
    """
    更新股票基本信息到数据库
    
    Args:
        stock_code_list (list, optional): 股票代码列表. 默认为None，表示更新所有股票信息
    
    Returns:
        bool: 操作是否成功
    """
    try:
        # 获取股票基本信息
        stock_df = fetch_stock_basic_info(stock_code_list)
        
        if stock_df.empty:
            logger.warning("未获取到股票基本信息数据，跳过更新")
            return False
            
        # 获取数据库会话
        db = SessionLocal()
        
        try:
            # 获取所有ts_code列表
            ts_code_list = stock_df['ts_code'].tolist()
            
            # 批量查询已存在的记录
            exist_stocks = db.query(StockBasicInfo).filter(
                StockBasicInfo.ts_code.in_(ts_code_list)
            ).all()
            
            # 构建已存在记录的映射字典
            exist_stock_map = {stock.ts_code: stock for stock in exist_stocks}
            
            # 用于统计更新和新增的数量
            update_count = 0
            insert_count = 0
            
            # 遍历股票数据，更新或插入
            for _, row in stock_df.iterrows():
                # 将数据转换为字典，并处理nan值
                stock_dict = row.where(pd.notnull(row), None).to_dict()
                ts_code = stock_dict['ts_code']
                
                if ts_code in exist_stock_map:
                    # 更新现有记录
                    exist_stock = exist_stock_map[ts_code]
                    for key, value in stock_dict.items():
                        if key in exist_stock.__table__.columns.keys() and key != 'id':
                            setattr(exist_stock, key, value)
                    update_count += 1
                else:
                    # 创建新记录
                    new_stock = StockBasicInfo(**stock_dict)
                    db.add(new_stock)
                    insert_count += 1
            
            # 提交事务
            db.commit()
            logger.info(f"成功更新股票基本信息：更新 {update_count} 条，新增 {insert_count} 条")
            return True
            
        except Exception as e:
            db.rollback()
            logger.error(f"更新股票基本信息到数据库失败: {e}")
            return False
        finally:
            db.close()
            
    except Exception as e:
        logger.error(f"更新股票基本信息过程异常: {e}")
        return False

if __name__ == "__main__":
    # 测试代码
    update_stock_basic_info() 